
/*
** application layer multicast client
** location: multicast-tree node
** use UDP && TCP
** gcc -lwsock32 to compile
*/

#include <stdio.h>
#include <stdlib.h>
#include <Winsock2.h>
#include <process.h>
#include <time.h>

#pragma comment(lib, "ws2_32.lib")

#define SERVER_PORT 12500
#define BUF_SIZE    1024
#define WSVERS      MAKEWORD(2, 2)
#define ROOTADDR    "127.0.0.1"
#define MAX_CHILD   2

struct Node {
    SOCKET sock;
    struct sockaddr_in addr;
};
struct Node childNode[MAX_CHILD];
struct Node parent;
int child_count = 0;

int testRecv(int sLen) {
    if (sLen == 0 || sLen == SOCKET_ERROR) {
        return -1;
    } else if (sLen > 0) {
        return 1;
    }
}

void testSend(int sLen) {
    switch (sLen) {
    case SOCKET_ERROR: printf("send() failed: %d\n", WSAGetLastError()); break;
    case 0: printf("send() failed: %d\n", WSAGetLastError()); break;
    default: break;
    }
    return ;
}

unsigned WINAPI sendmsg(void * arg) {
    printf("send Hi~\n");
    while (true) {
        // 
    }
}
unsigned WINAPI recvmsg(void * arg) {
    printf("recv Hi~\n");
    while (true) {
        //
    }
}

int main(int argc, char *argv[]) {
    struct sockaddr_in servAddr, clntAddr;
    SOCKET sock;
    int addrLen = sizeof(clntAddr);
    char wBuf[BUF_SIZE];
    char rBuf[BUF_SIZE];
    int sLen;
    HANDLE   hSendThread, hRecvThread;
    WSADATA  wsaData;
    if (WSAStartup(WSVERS, &wsaData) != 0)
        printf("WSAStartup() error! \n");

    sock = socket(PF_INET, SOCK_STREAM, 0);
    if (sock == INVALID_SOCKET)
        printf("TCP socket creation error! \n");

    // server info
    memset(&servAddr, 0, sizeof(servAddr));
    servAddr.sin_family = AF_INET;
    servAddr.sin_addr.s_addr = inet_addr(argv[1]);//(ROOTADDR);
    servAddr.sin_port = htons(SERVER_PORT);
    
    // connect to RP for register
    if (connect(sock, (SOCKADDR*)&servAddr, addrLen) == INVALID_SOCKET) {
        printf("connect() failed: %d\n.", WSAGetLastError());
    } else {
        printf("connect to %s:%d \n", inet_ntoa(servAddr.sin_addr), ntohs(servAddr.sin_port));
    }
    // get port
    char port[20];
    memset(rBuf, 0, sizeof(rBuf));
    sLen = recv(sock, rBuf, sizeof(rBuf), 0);
    if (testRecv(sLen) < 0) {
        printf("port recv error.\n");
    } else {
        printf("%s\n", rBuf);
        strncpy(port, rBuf, strlen(rBuf) + 1);
    }

    closesocket(sock);
    // be a server to wait for connection
    sock = socket(PF_INET, SOCK_STREAM, 0);
    if (sock == INVALID_SOCKET) {
        printf("socket created error.\n");
    }
    memset(&servAddr, 0, sizeof(servAddr));
    servAddr.sin_family = AF_INET;
    servAddr.sin_addr.s_addr = htonl(INADDR_ANY);
    servAddr.sin_port = htons(atoi(port));
    // binding
    if (bind(sock, (SOCKADDR*)&servAddr, addrLen) == SOCKET_ERROR) {
        printf("bind() failed: %d\n", WSAGetLastError());
    }
    // listen
    if (listen(sock, 0) == SOCKET_ERROR) {
        printf("listen() failed: %d\n", WSAGetLastError());
    }

    // get parent node
    parent.sock = accept(sock, (SOCKADDR*)&parent.addr, &addrLen);
    if (parent.sock == INVALID_SOCKET) {
        printf("accept() failed: %d\n", WSAGetLastError());
    } else {
        // inform parent
        memset(rBuf, 0, sizeof(rBuf));
        sLen = recv(parent.sock, rBuf, strlen(rBuf), 0);
        printf("%s be a child node of %s:%d \n", rBuf,
                inet_ntoa(parent.addr.sin_addr), ntohs(parent.addr.sin_port));
    }

    // recv multicast message from parent node
    while (true) {
        printf("wait for multicast message...\n");
        memset(rBuf, 0, sizeof(rBuf));
        sLen = recv(parent.sock, rBuf, sizeof(rBuf), 0);
        if (testRecv(sLen) < 0) {
            printf("multicast message recv error: %d\n", WSAGetLastError());
        } else {

            // new_node
            // request to be a child node
            if (strncmp(rBuf, "new_node", 9) == 0) {
                if (child_count < MAX_CHILD) {
                    printf("%s\n", rBuf);

                    // recv buffer
                    char ip_temp[20];
                    char port_temp[20];

                    // IP and port
                    memset(rBuf, 0, sizeof(rBuf));
                    sLen = recv(parent.sock, rBuf, sizeof(rBuf), 0);
                    if (testRecv(sLen) < 0) {
                        printf("new node info recv error.\n");
                    } else {

                        char *p;
                        p = strtok(rBuf, ",");
                        if(p) {
                            printf("%s,", p);
                            strncpy(ip_temp, p, strlen(p) + 1);
                        }
                        p = strtok(NULL,",");
                        if(p) {
                            printf("%s\n", p);
                            strncpy(port_temp, p, strlen(p) + 1);
                        }

                    }

                    if (child_count < MAX_CHILD) {

                        // make it be a directed child
                        // create socket
                        childNode[child_count].sock = socket(PF_INET, SOCK_STREAM, 0);
                        if (childNode[child_count].sock == INVALID_SOCKET)
                            printf("socket created error.\n");
                        // addr info
                        childNode[child_count].addr.sin_family = AF_INET;
                        childNode[child_count].addr.sin_addr.s_addr = inet_addr(ip_temp);
                        childNode[child_count].addr.sin_port = htons(atoi(port_temp));
                        // connect
                        SOCKET dest = connect(childNode[child_count].sock, 
                            (SOCKADDR*)&childNode[child_count].addr, addrLen);
                        if (dest == INVALID_SOCKET) {
                            // may be connected to other parent node
                            printf("new node may be connected to other node: %d \n", WSAGetLastError());

                            continue;
                        } else {
                            
                            // connect successfully
                            memset(wBuf, 0, sizeof(wBuf));
                            sprintf(wBuf, "success");
                            sLen = send(childNode[child_count].sock, wBuf, sizeof(wBuf), 0);
                            if (sLen == 0 || sLen == SOCKET_ERROR) {
                                printf("connect() failed: %d\n", WSAGetLastError());
                            } else {
                                printf("%s:%d become a child node.\n",
                                    inet_ntoa(childNode[child_count].addr.sin_addr), 
                                    ntohs(childNode[child_count].addr.sin_port));

                                child_count++;
                            }
                            continue;
                        }
                    }

                } else {
                    srand(time(0));
                    int i = rand() % child_count;

                    sLen = send(childNode[i].sock, rBuf, strlen(rBuf), 0);
                    testSend(sLen);
                    memset(rBuf, 0, sizeof(rBuf));
                    sLen = recv(parent.sock, rBuf, sizeof(rBuf), 0);
                    if (testRecv(sLen) < 0) {
                        printf("addr info recv error.\n");
                    } else {
                        sLen = send(childNode[i].sock, rBuf, strlen(rBuf), 0);
                        testSend(sLen);
                    }
                }

            } else {
                // transfer to child nodes
                for (int i = 0; i < child_count; ++i) {
                    //printf("transfer to %s:%d\n", inet_ntoa(childNode[i].addr.sin_addr), 
                    //    ntohs(childNode[i].addr.sin_port));
                    sLen = send(childNode[i].sock, rBuf, strlen(rBuf), 0);
                    testSend(sLen);
                }
            }
            
            // normal multicast message
            printf("%s", rBuf);
        }
    }

/*
    hSendThread = (HANDLE)_beginthreadex(NULL, 0, &sendmsg, NULL, 0, NULL);
    hRecvThread = (HANDLE)_beginthreadex(NULL, 0, &recvmsg, NULL, 0, NULL);

    // wait recv thread and send thread
    WaitForSingleObject(hSendThread, INFINITE);
    WaitForSingleObject(hRecvThread, INFINITE);
*/

    return 0;
}

